﻿using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Net;
using System.Threading;
using System.Threading.Tasks;
using Amazon.S3;
using Amazon.S3.Model;
using Polly;

public class S3BucketFileDownloader
{
    private readonly IAmazonS3 _s3Client;
    private const int MaxRetries = 3;
    private readonly TimeSpan RetryDelay = TimeSpan.FromSeconds(2);
    private readonly SemaphoreSlim _downloadSemaphore = new(4, 4); // 最大并发下载数4

    private const int DownloadBufferSize = 1 * 1024 * 1024; // 50KB缓冲区
    private const long DownloadMultipartThreshold = 100 * 1024 * 1024; // 100MB
    private const int DownloadPartSizeMB = 20 * 1024 * 1024; // 分段大小20MB

    // 下载进度事件
    public event Action<ProgressReport> ProgressChanged;

    private DownloadTracker _downloadTracker;
    private MergeTracker _mergeTracker;
    private BlockingCollection<CompletedPart> _completedPartsQueue;
    private CancellationTokenSource _cts;
    private Task _mergeTask;

    public S3BucketFileDownloader(IAmazonS3 s3Client)
    {
        _s3Client = s3Client;
    }

    public async Task DownloadFileAsync(string bucketName, string s3Key, string localFilePath)
    {
        _cts = new CancellationTokenSource();
        _completedPartsQueue = new BlockingCollection<CompletedPart>();

        try
        {
            var metadata = await GetObjectMetadata(bucketName, s3Key);
            var fileSize = metadata.ContentLength;

            if (fileSize > DownloadMultipartThreshold)
            {
                await DownloadLargeFile(bucketName, s3Key, localFilePath, fileSize);
            }
            else
            {
                await DownloadSmallFile(bucketName, s3Key, localFilePath, fileSize);
            }
        }
        catch (Exception ex)
        {
            _cts.Cancel();

            if (File.Exists(localFilePath))
                File.Delete(localFilePath);

            throw new Exception("文件下载失败", ex);
        }
        finally
        {
            _completedPartsQueue?.Dispose();
            _cts?.Dispose();
        }
    }

    private async Task<GetObjectMetadataResponse> GetObjectMetadata(string bucketName, string key)
    {
        return await _s3Client.GetObjectMetadataAsync(new GetObjectMetadataRequest
        {
            BucketName = bucketName,
            Key = key
        });
    }

    private async Task DownloadSmallFile(string bucketName, string s3Key, string localFilePath, long fileSize)
    {
        _downloadTracker = new DownloadTracker(fileSize, 1); // 小文件视为1个部分

        using var progressTimer = new Timer(ProgressReportCallback, null,
            ProgressReport.ReportTimeSpan, ProgressReport.ReportTimeSpan);

        using var response = await _s3Client.GetObjectAsync(new GetObjectRequest
        {
            BucketName = bucketName,
            Key = s3Key
        });

        await using var fileStream = new FileStream(
            localFilePath,
            FileMode.Create,
            FileAccess.Write,
            FileShare.None,
            bufferSize: DownloadBufferSize,
            useAsync: true);

        var buffer = new byte[DownloadBufferSize];
        long totalRead = 0;

        while (true)
        {
            var bytesRead = await response.ResponseStream.ReadAsync(buffer);
            if (bytesRead == 0) break;

            await fileStream.WriteAsync(buffer.AsMemory(0, bytesRead));
            totalRead += bytesRead;
            _downloadTracker.UpdateProgress(1, totalRead, fileSize);
        }

        _downloadTracker.MarkComplete(1);
        ProgressReportCallback(null); // 最终报告
    }

    private async Task DownloadLargeFile(string bucketName, string s3Key, string localFilePath, long fileSize)
    {
        long partSize = DownloadPartSizeMB;
        var partCount = (int)Math.Ceiling((double)fileSize / partSize);

        // 初始化跟踪器
        _downloadTracker = new DownloadTracker(fileSize, partCount);
        _mergeTracker = new MergeTracker(fileSize, partCount);

        // 启动合并线程
        _mergeTask = Task.Run(() => MergeFilesAsync(localFilePath, _cts.Token));

        // 启动进度报告
        using var progressTimer = new Timer(ProgressReportCallback, null,
            ProgressReport.ReportTimeSpan, ProgressReport.ReportTimeSpan);

        // 创建并执行所有分片下载任务
        var downloadTasks = Enumerable.Range(1, partCount).Select(async partNumber =>
        {
            var start = (partNumber - 1) * partSize;
            var end = Math.Min(partNumber * partSize - 1, fileSize - 1);
            var partSizeActual = end - start + 1;
            var tempFile = $"{localFilePath}.part{partNumber}";

            await _downloadSemaphore.WaitAsync(_cts.Token);
            try
            {
                await DownloadPartWithRetryAsync(
                    bucketName, s3Key, tempFile,
                    start, end, partNumber, partSizeActual);

                // 将完成的文件加入队列(带partNumber)
                _completedPartsQueue.Add(new CompletedPart
                {
                    PartNumber = partNumber,
                    FilePath = tempFile
                }, _cts.Token);
            }
            finally
            {
                _downloadSemaphore.Release();
            }
        });

        await Task.WhenAll(downloadTasks);
        _completedPartsQueue.CompleteAdding(); // 标记下载完成

        await _mergeTask; // 等待合并完成
        ProgressReportCallback(null); // 最终报告
    }

    private async Task DownloadPartWithRetryAsync(
        string bucketName,
        string s3Key,
        string partFilePath,
        long startByte,
        long endByte,
        int partNumber,
        long partSize)
    {
        var policy = Policy
            .Handle<AmazonS3Exception>(IsRetryableError)
            .Or<IOException>()
            .WaitAndRetryAsync(MaxRetries, _ => RetryDelay);

        await policy.ExecuteAsync(async () =>
        {
            using var response = await _s3Client.GetObjectAsync(new GetObjectRequest
            {
                BucketName = bucketName,
                Key = s3Key,
                ByteRange = new ByteRange(startByte, endByte)
            });

            await using var fileStream = new FileStream(
                partFilePath,
                FileMode.Create,
                FileAccess.Write,
                FileShare.None,
                bufferSize: DownloadBufferSize,
                useAsync: true);

            var buffer = new byte[DownloadBufferSize];
            long partDownloaded = 0;

            while (true)
            {
                var bytesRead = await response.ResponseStream.ReadAsync(buffer);
                if (bytesRead == 0) break;

                await fileStream.WriteAsync(buffer.AsMemory(0, bytesRead));
                partDownloaded += bytesRead;
                _downloadTracker.UpdateProgress(partNumber, partDownloaded, partSize);
            }

            _downloadTracker.MarkComplete(partNumber);
        });
    }

    private async Task MergeFilesAsync(string outputFilePath, CancellationToken ct)
    {
        await using var outputStream = new FileStream(
            outputFilePath,
            FileMode.Create,
            FileAccess.Write,
            FileShare.None,
            bufferSize: DownloadBufferSize,
            useAsync: true);

        // 使用SortedDictionary来按partNumber排序
        var pendingParts = new SortedDictionary<int, string>();
        int nextExpectedPart = 1; // 从第1部分开始

        foreach (var completedPart in _completedPartsQueue.GetConsumingEnumerable(ct))
        {
            // 将收到的部分添加到待处理字典
            pendingParts.Add(completedPart.PartNumber, completedPart.FilePath);

            // 检查是否可以合并下一个部分
            while (pendingParts.TryGetValue(nextExpectedPart, out var fileToMerge))
            {
                pendingParts.Remove(nextExpectedPart);
                await MergeSinglePart(outputStream, fileToMerge, nextExpectedPart);
                nextExpectedPart++;
            }
        }
    }

    private async Task MergeSinglePart(FileStream outputStream, string partFile, int partNumber)
    {
        var fileInfo = new FileInfo(partFile);
        long partMerged = 0;

        try
        {
            await using var partStream = new FileStream(
                partFile,
                FileMode.Open,
                FileAccess.Read,
                FileShare.Read,
                bufferSize: DownloadBufferSize,
                useAsync: true);

            var buffer = new byte[DownloadBufferSize];
            int bytesRead;

            while ((bytesRead = await partStream.ReadAsync(buffer)) > 0)
            {
                await outputStream.WriteAsync(buffer.AsMemory(0, bytesRead));
                partMerged += bytesRead;
                _mergeTracker.UpdateProgress(partNumber, bytesRead, fileInfo.Length);
            }

            _mergeTracker.MarkComplete(partNumber);
        }
        finally
        {
            try { File.Delete(partFile); } catch { /* 忽略删除错误 */ }
        }
    }

    private void ProgressReportCallback(object state)
    {
        ProgressChanged?.Invoke(new ProgressReport
        {
            Download = _downloadTracker,
            Merge = _mergeTracker,
            ReportTime = DateTime.Now
        });
    }

    private bool IsRetryableError(AmazonS3Exception ex)
    {
        return ex.ErrorCode == "RequestTimeout" ||
               ex.StatusCode == HttpStatusCode.ServiceUnavailable ||
               ex.StatusCode == HttpStatusCode.InternalServerError;
    }

    private class CompletedPart
    {
        public int PartNumber { get; set; }
        public string FilePath { get; set; }
    }
}

public class DownloadTracker
{
    private class PartProgress
    {
        public long DownloadedBytes { get; set; }
        public long TotalBytes { get; set; }
    }

    public List<int> AllParts { get; } // 全部分包列表
    public ConcurrentBag<int> CompletedParts { get; } = new(); // 已下载分包列表
    public long TotalBytes { get; } // 总下载字节数
    public long DownloadedBytes => _partProgress.Values.Sum(p => p.DownloadedBytes);
    public double Percentage => Math.Round(DownloadedBytes * 100d / TotalBytes, 2);
    public bool IsCompleted => CompletedParts.Count == AllParts.Count;
    public double DownloadSpeedMBps { get; private set; }

    private readonly ConcurrentDictionary<int, PartProgress> _partProgress = new();
    private readonly object _speedLock = new();
    private DateTime _lastUpdate;
    private long _lastTotalDownloaded;

    public DownloadTracker(long totalBytes, int partCount)
    {
        TotalBytes = totalBytes;
        AllParts = Enumerable.Range(1, partCount).ToList();
        _lastUpdate = DateTime.Now;
        _lastTotalDownloaded = 0;
    }

    public void UpdateProgress(int partNumber, long partDownloaded, long partSize)
    {
        _partProgress.AddOrUpdate(partNumber,
            _ => new PartProgress { DownloadedBytes = partDownloaded, TotalBytes = partSize },
            (_, existing) =>
            {
                existing.DownloadedBytes = partDownloaded;
                return existing;
            });

        lock (_speedLock)
        {
            var now = DateTime.Now;
            var timeSinceLast = now - _lastUpdate;

            if (timeSinceLast.TotalSeconds > 0)
            {
                var currentTotal = DownloadedBytes;
                var bytesSinceLast = currentTotal - _lastTotalDownloaded;
                DownloadSpeedMBps = (bytesSinceLast / 1024.0 / 1024.0) / timeSinceLast.TotalSeconds;

                _lastTotalDownloaded = currentTotal;
                _lastUpdate = now;
            }
        }
    }

    public void MarkComplete(int partNumber)
    {
        CompletedParts.Add(partNumber);
    }
}

public class MergeTracker
{
    public int TotalParts { get; }
    public int MergedPartsCount => _mergedParts.Count;
    public long TotalBytes { get; private set; }
    public long MergedBytes => _mergedBytes;
    public double Percentage => TotalBytes > 0 ? Math.Round(MergedBytes * 100d / TotalBytes, 2) : 0;
    public double MergeSpeedMBps { get; private set; }
    public bool IsCompleted => MergedPartsCount == TotalParts;

    private long _mergedBytes;
    private readonly ConcurrentBag<int> _mergedParts = new();
    private readonly object _speedLock = new();
    private DateTime _lastUpdate;
    private long _lastTotalMerged;

    public MergeTracker(long totalBytes, int totalParts)
    {
        TotalBytes = totalBytes;
        TotalParts = totalParts;
        _lastUpdate = DateTime.Now;
        _lastTotalMerged = 0;
    }

    public void UpdateProgress(int partNumber, long bytesMerged, long fileSize)
    {
        Interlocked.Add(ref _mergedBytes, bytesMerged);
        //TotalBytes += fileSize;

        lock (_speedLock)
        {
            var now = DateTime.Now;
            var timeSinceLast = now - _lastUpdate;

            if (timeSinceLast.TotalSeconds > 0)
            {
                var currentTotal = MergedBytes;
                var bytesSinceLast = currentTotal - _lastTotalMerged;
                MergeSpeedMBps = (bytesSinceLast / 1024.0 / 1024.0) / timeSinceLast.TotalSeconds;

                _lastTotalMerged = currentTotal;
                _lastUpdate = now;
            }
        }
    }

    public void MarkComplete(int partNumber)
    {
        _mergedParts.Add(partNumber);
    }
}

public class ProgressReport
{
    public static TimeSpan ReportTimeSpan = TimeSpan.FromMilliseconds(1000);
    public DownloadTracker Download { get; set; }
    public MergeTracker Merge { get; set; }
    public DateTime ReportTime { get; set; } = DateTime.Now;
}